from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64

class CryptoUtils:
    """
    加密工具类
    提供了一系列静态方法用于加密通信和密码处理：
    - RSA非对称加密（密钥生成、序列化、加密解密）
    - Fernet对称加密（密钥生成、消息加密解密）
    - 密码哈希处理
    """

    @staticmethod
    def generate_key_pair():
        """
        1.RSA非对称加密：生成RSA密钥对（服务器启动时执行一次）
        Returns:
            tuple: (私钥, 公钥)
        """
        # 生成2048位的RSA密钥对
        private_key = rsa.generate_private_key(
            public_exponent=65537, # 常用的公钥指数
            key_size=2048   # 密钥长度
        )
        public_key = private_key.public_key()
        return private_key, public_key

    @staticmethod
    def serialize_public_key(public_key):
        """
        2.序列化公钥为PEM格式（使用Base64编码的二进制数据）
        Args:
            public_key: RSA公钥对象
        Returns:
            bytes: PEM格式的公钥数据
        """
        # 将公钥序列化为PEM格式
        pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        return pem

    @staticmethod
    def deserialize_public_key(pem):
        """
        3.从PEM格式数据反序列化公钥
        Args:
            pem (bytes): PEM格式的公钥数据
        Returns:
            RSA公钥对象
        """
        return serialization.load_pem_public_key(pem)



    @staticmethod
    def encrypt_with_public_key(public_key, message):
        """
        使用RSA公钥加密数据
        Args:
            public_key: RSA公钥对象
            message (bytes): 待加密数据
        Returns:
            bytes: 加密后的数据
        """
        encrypted = public_key.encrypt(
            message,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return encrypted

    @staticmethod
    def decrypt_with_private_key(private_key, encrypted_message):
        """
        使用RSA私钥解密数据
        Args:
            private_key: RSA私钥对象
            encrypted_message (bytes): 加密数据
        Returns:
            bytes: 解密后的数据
        """
        decrypted = private_key.decrypt(
            encrypted_message,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return decrypted

    @staticmethod
    def generate_symmetric_key():
        """
        生成Fernet对称加密密钥
        Returns:
            bytes: 对称加密密钥
        """
        return Fernet.generate_key()

    @staticmethod
    def encrypt_message(key, message):
        """
        使用Fernet对称密钥加密消息
        Args:
            key (bytes): Fernet密钥
            message (str): 待加密消息
        Returns:
            bytes: 加密后的数据
        """
        f = Fernet(key)
        return f.encrypt(message.encode())

    @staticmethod
    def decrypt_message(key, encrypted_message):
        """
        使用Fernet对称密钥解密消息
        Args:
            key (bytes): Fernet密钥
            encrypted_message (bytes): 加密数据
        Returns:
            str: 解密后的消息
        """
        f = Fernet(key)
        return f.decrypt(encrypted_message).decode()

    @staticmethod
    def hash_password(password):
        """
        使用PBKDF2进行密码哈希
        Args:
            password (str): 原始密码
        Returns:
            bytes: base64编码的密码哈希值
        """
        salt = b'salt_123'  # salt值
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        return base64.b64encode(kdf.derive(password.encode()))
